﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Diagnostics;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Runtime.Serialization.Formatters.Binary;
using System.Security;
using System.Text;
using System.Windows.Forms;
using System.Xml.Linq;
using OpenCLNet;
using Image = OpenCLNet.Image;

namespace GPUTracer
{
    public partial class TracerForm : Form
    {
        private bool Initialized = false;
        private Bitmap OutputBitmap = null;
        private OpenCLManager OCLMan;
        private OpenCLNet.Program OCLProgram;
        private Kernel RayTraceKernel;

        private Float4 cameraPosition;
        private Float4 cameraLookAt;
        private Float4 cameraForward;
        private Mem _lightsBuffer;
        private Sampler _sampler;
        private Stopwatch _drawingStopwatch = new Stopwatch();
        private Stopwatch _lastRenderedFrame = new Stopwatch();

        private XElement _sceneXML = XElement.Load("Scene.xml");
        private Image _skyTexture;
        private Image _woodTexture;
        private List<Object> _objects;
        private KDTreeElement _tree;
        private int _numPlanes;
        private List<Light> _lights;

        private const int Blocksize = 512;

        private const int Depth = 3;
        private const int ShortStackSize = 3;
        private const bool UseKDTree = true;
        private const bool ShowShadows = true;
        private const bool GPURayTracing = true;

        public TracerForm()
        {
            InitializeComponent();
            if (!GPURayTracing)
                Text = "CPUTracer";
        }

        private void Draw()
        {
            if (GPURayTracing)
            {
                BitmapData bd = OutputBitmap.LockBits(new Rectangle(0, 0, OutputBitmap.Width, OutputBitmap.Height),
                                                      ImageLockMode.WriteOnly, OutputBitmap.PixelFormat);
                DoRayTracing(bd);
                OutputBitmap.UnlockBits(bd);
            }
            else
            {
                Float4 right;
                Float4 up;
                GetCameraParameters(out right, out up);
                CPURayTracer.RayTrace(OutputBitmap, _tree, _objects.ToArray(), _numPlanes, _lights.ToArray(), cameraPosition, cameraForward, right, up);
            }
        }

        private void DoRayTracing(BitmapData output)
        {
            int deviceIndex = comboBoxDeviceSelector.SelectedIndex;
            Mem outputBuffer = OCLMan.Context.CreateBuffer(MemFlags.USE_HOST_PTR, output.Stride * output.Height, output.Scan0);

            Float4 right;
            Float4 up;
            GetCameraParameters(out right, out up);

            for (int x = 0; x < output.Width; x += Blocksize)
            {
                for (int y = 0; y < output.Height; y += Blocksize)
                {
                    var rayTracingGlobalWorkSize = new IntPtr[2];
                    rayTracingGlobalWorkSize[0] = (IntPtr)(output.Width - x > Blocksize ? Blocksize : output.Width - x);
                    rayTracingGlobalWorkSize[1] = (IntPtr)(output.Height - y > Blocksize ? Blocksize : output.Height - y);
                    var rayTracingGlobalOffset = new IntPtr[2];
                    rayTracingGlobalOffset[0] = (IntPtr)x;
                    rayTracingGlobalOffset[1] = (IntPtr)y;

                    RayTraceKernel.SetArg(0, output.Width);
                    RayTraceKernel.SetArg(1, output.Height);
                    RayTraceKernel.SetArg(2, outputBuffer);
                    RayTraceKernel.SetArg(3, output.Stride);
                    RayTraceKernel.SetArg(4, cameraPosition);
                    RayTraceKernel.SetArg(5, cameraForward);
                    RayTraceKernel.SetArg(6, right);
                    RayTraceKernel.SetArg(7, up);
                    RayTraceKernel.SetArg(8, _skyTexture);
                    RayTraceKernel.SetArg(9, _woodTexture);
                    RayTraceKernel.SetArg(10, _sampler);

                    OCLMan.CQ[deviceIndex].EnqueueNDRangeKernel(RayTraceKernel, 2, rayTracingGlobalOffset, rayTracingGlobalWorkSize, null);
                }
            }
            
            OCLMan.CQ[deviceIndex].EnqueueBarrier();
            IntPtr p = OCLMan.CQ[deviceIndex].EnqueueMapBuffer(outputBuffer, true, MapFlags.READ, IntPtr.Zero, (IntPtr)(output.Stride * output.Height));
            OCLMan.CQ[deviceIndex].EnqueueUnmapMemObject(outputBuffer, p);
            OCLMan.CQ[deviceIndex].Finish();

            outputBuffer.Dispose();
        }

        private void GetCameraParameters(out Float4 right, out Float4 up)
        {
            Float4 down = new Float4(0.0f, -1.0f, 0.0f, 0.0f);
            right = Helper.Cross(cameraForward, down).Normalize().Times(1.5f);
            up = Helper.Cross(cameraForward, right).Normalize().Times(-1.5f);
        }

        private void InitCamera()
        {
            if (_sceneXML.Element("camera") != null)
            {
                cameraPosition = GetFloat4FromXElement(_sceneXML.Element("camera").Element("position"));
                cameraLookAt = GetFloat4FromXElement(_sceneXML.Element("camera").Element("lookAt"));
                cameraForward = cameraLookAt.Sub(cameraPosition).Normalize();
            }
            else
            {
                throw new Exception("No camera informations in scene.xml!");
            }
        }

        private byte[] GetScene()
        {
            if (GPURayTracing && File.Exists("tree.bin") && (File.GetLastWriteTime("scene.xml") < File.GetLastWriteTime("tree.bin")))
            {
                var rstream = File.Open("tree.bin", FileMode.Open);
                byte[] l = new byte[4];
                rstream.Read(l, 0, 4);
                int count = BitConverter.ToInt32(l, 0);
                byte[] res = new byte[count];
                rstream.Read(res, 0, count);
                return res;
            }

            _objects = new List<Object>();
            _numPlanes = 0;
            foreach (var element in _sceneXML.Elements())
            {
                try
                {
                    switch (element.Name.ToString())
                    {
                        case "sphere":
                            var r = float.Parse(element.Attribute("radius").Value);
                            var pos = GetFloat4FromXElement(element.Element("position"));
                            if (element.Attribute("surface") != null)
                            {
                                var surface = GetSurfaceTypeFromString(element.Attribute("surface").Value);
                                _objects.Add(new Sphere(pos, r, surface));
                            }
                            else
                                _objects.Add(new Sphere(pos, r));
                            break;
                        case "plane":
                            var offset = float.Parse(element.Attribute("offset").Value);
                            var normal = GetFloat4FromXElement(element.Element("normal"));
                            if (element.Attribute("surface") != null)
                            {
                                var surface = GetSurfaceTypeFromString(element.Attribute("surface").Value);
                                _objects.Insert(0, new Plane(normal, offset, surface));
                            }
                            else
                                _objects.Insert(0, new Plane(normal, offset));
                            _numPlanes++;
                            break;
                        case "triangle":
                            var p1 = GetFloat4FromXElement(element.Element("p1"));
                            var p2 = GetFloat4FromXElement(element.Element("p2"));
                            var p3 = GetFloat4FromXElement(element.Element("p3"));
                            if (element.Attribute("surface") != null)
                            {
                                var surface = GetSurfaceTypeFromString(element.Attribute("surface").Value);
                                _objects.Add(new Triangle(p1, p2, p3, surface));
                            }
                            else
                                _objects.Add(new Triangle(p1, p2, p3));
                            break;
                    }
                }
                catch (Exception)
                {
                }
            }

            _tree = KDTreeBuilder.Build(_objects.ToArray());
            var tw = new System.IO.StreamWriter("tree.txt");
            tw.Write(_tree);
            tw.Close();

            var kdTreeConverter = new KDTreeConverter();
            var result = kdTreeConverter.ConvertKDTree(_tree, _objects.ToArray());

            var stream = File.Open("tree.bin", FileMode.Create);
            stream.Write(BitConverter.GetBytes(result.Length), 0, 4);
            stream.Write(result, 0, result.Length);
            stream.Close();

            return result;
        }

        private Object.SurfaceType GetSurfaceTypeFromString(string value)
        {
            switch (value.ToLower())
            {
                case "checkerboard":
                    return Object.SurfaceType.Checkerboard;
                case "shiny":
                    return Object.SurfaceType.Shiny;
                case "transparent":
                    return Object.SurfaceType.Transparent;
                case "sky":
                    return Object.SurfaceType.Sky;
                case "wood":
                    return Object.SurfaceType.Wood;
                case "transparentandshiny":
                default:
                    return Object.SurfaceType.TransparentAndShiny;
            }
        }

        private Float4 GetFloat4FromXElement(XElement element)
        {
            var s0 = float.Parse(element.Attribute("S0").Value);
            var s1 = float.Parse(element.Attribute("S1").Value);
            var s2 = float.Parse(element.Attribute("S2").Value);
            return new Float4(s0, s1, s2, 0);
        }

        private Float4 GetFloat4FromXElement2(XElement element)
        {
            var s0 = float.Parse(element.Attribute("R").Value);
            var s1 = float.Parse(element.Attribute("G").Value);
            var s2 = float.Parse(element.Attribute("B").Value);
            return new Float4(s0, s1, s2, 0);
        }

        private byte[] GetLights()
        {
            _lights = new List<Light>();
            foreach (var element in _sceneXML.Elements())
            {
                try
                {
                    if (element.Name.ToString() == "light")
                    {
                        var pos = GetFloat4FromXElement(element.Element("position"));
                        var color = GetFloat4FromXElement2(element.Element("color"));
                        _lights.Add(new Light(pos, color));
                    }
                }
                catch (Exception)
                {
                }
            }

            var memstream = new MemoryStream();
            var sceneStream = new BinaryWriter(memstream);

            var numHeader = new byte[4 * 4];
            Buffer.BlockCopy(BitConverter.GetBytes(_lights.Count), 0, numHeader, 0, 4);

            sceneStream.Write(numHeader);

            foreach (var light in _lights)
            {
                sceneStream.Write(light.ToBinary());
            }

            return memstream.ToArray();
        }

        private void TracerForm_Load(object sender, EventArgs e)
        {
            try
            {
                InitializeOpenCL();
                Initialized = true;
            }
            catch (Exception ex)
            {
                MessageBox.Show(ex.Message, "Initiation failed");
                Application.Exit();
            }
        }

        private void InitializeOpenCL()
        {
            if (GPURayTracing)
            {
                if (OpenCL.NumberOfPlatforms == 0)
                {
                    MessageBox.Show("OpenCL not available!");
                    Application.Exit();
                }

                OCLMan = new OpenCLManager();
                OCLMan.AttemptUseBinaries = true;
                OCLMan.AttemptUseSource = true;
                OCLMan.RequireImageSupport = true;

                OCLMan.Defines = string.Format("#define _DEPTH {0}\n", Depth);
                OCLMan.Defines += string.Format("#define _SHORT_STACK_SIZE {0}\n", ShortStackSize);
                if (!UseKDTree)
                    OCLMan.Defines += "#define _BRUTEFORCERENDERING_\n";
                if (ShowShadows)
                    OCLMan.Defines += "#define _SHOW_SHADOWS_";
                //OCLMan.BuildOptions = "-cl-opt-disable";

                OCLMan.CreateDefaultContext(0, DeviceType.ALL);
                var rtCode =
                    Properties.Resources.RayTracer.Replace("$$SCENEDATA$$", GetSceneString()).Replace("$$LIGHTSDATA$$",
                                                                                                      GetLightsString());
                OCLProgram = OCLMan.CompileSource(rtCode);

                for (int i = 0; i < OCLMan.Context.Devices.Length; i++)
                    comboBoxDeviceSelector.Items.Add(OCLMan.Context.Devices[i].Vendor + ":" +
                                                     OCLMan.Context.Devices[i].Name);
                comboBoxDeviceSelector.SelectedIndex = 0;

                RayTraceKernel = OCLProgram.CreateKernel("RayTrace");

                ImageLoader.OclContext = OCLMan.Context;
                _sampler = OCLMan.Context.CreateSampler(true, AddressingMode.REPEAT, FilterMode.LINEAR);

                InitCamera();

                _skyTexture = ImageLoader.LoadImage("skyTexture.png");
                _woodTexture = ImageLoader.LoadImage("woodTexture.JPG");
            }
            else
            {
                InitCamera();
                GetScene();
                GetLights();
            }
        }

        private string GetSceneString()
        {
            byte[] sceneData = GetScene();
            var sb = new StringBuilder();
            foreach (byte b in sceneData)
            {
                sb.Append(String.Format("0x{0:X2}, ", b));
            }
            sb.Remove(sb.Length - 2, 2);
            return sb.ToString();
        }

        private string GetLightsString()
        {
            byte[] lightsData = GetLights();
            var sb = new StringBuilder();
            foreach (byte b in lightsData)
            {
                sb.Append(String.Format("0x{0:X2}, ", b));
            }
            sb.Remove(sb.Length - 2, 2);
            return sb.ToString();
        }

        private void DrawPanel_Paint(object sender, PaintEventArgs e)
        {
            _drawingStopwatch.Restart();

            if (OutputBitmap == null || (OutputBitmap.Width != DrawPanel.Width && OutputBitmap.Height != DrawPanel.Height))
                OutputBitmap = new Bitmap(DrawPanel.Width, DrawPanel.Height, PixelFormat.Format32bppArgb);

            Graphics g = e.Graphics;

            //g.FillRectangle(Brushes.Black, 0, 0, DrawPanel.Width, DrawPanel.Height);
            if (Initialized)
            {
                Draw();
                g.DrawImageUnscaled(OutputBitmap, 0, 0);
                //g.DrawImage(OutputBitmap, 0, 0, 1000, 1000);
            }


            _drawingStopwatch.Stop();
            var fps = 1000.0/_drawingStopwatch.ElapsedMilliseconds;
            string title;
            if (GPURayTracing)
                title = "GPUTracer";
            else
                title = "CPUTracer";
            Text = string.Format("{2} - FPS: {1} - Drawing Image Time: {0} ms", 
                _drawingStopwatch.ElapsedMilliseconds, fps, title);
        }

        private void Redraw()
        {
            DrawPanel.Refresh();
        }

        public void Idle()
        {
            short state = NativeMethods.GetKeyState((int) Keys.Up);
            bool upKeyDown = NativeMethods.HIWORD(state) != 0;
            state = NativeMethods.GetKeyState((int)Keys.Down);
            bool downKeyDown = NativeMethods.HIWORD(state) != 0;
            state = NativeMethods.GetKeyState((int)Keys.Left);
            bool leftKeyDown = NativeMethods.HIWORD(state) != 0;
            state = NativeMethods.GetKeyState((int)Keys.Right);
            bool rightKeyDown = NativeMethods.HIWORD(state) != 0;
            state = NativeMethods.GetKeyState((int)Keys.PageDown);
            bool pageDownKeyDown = NativeMethods.HIWORD(state) != 0;
            state = NativeMethods.GetKeyState((int)Keys.PageUp);
            bool pageUpKeyDown = NativeMethods.HIWORD(state) != 0;

            float scale = (_lastRenderedFrame.ElapsedMilliseconds) / 300.0f;
            _lastRenderedFrame.Restart();
            Float4 down = new Float4(0.0f, -1.0f, 0.0f, 0.0f);
            if (upKeyDown)
            {
                cameraPosition = cameraPosition.Add(new Float4(cameraForward.S0, 0.0f, cameraForward.S2, 0.0f).Times(scale));
            }
            if (downKeyDown)
            {
                cameraPosition =
                    cameraPosition.Sub(new Float4(cameraForward.S0, 0.0f, cameraForward.S2, 0.0f).Times(scale));
            }
            if (leftKeyDown)
            {
                cameraForward =
                    cameraForward.Add(Helper.Cross(down, cameraForward).Normalize().Times(0.1f).Times(scale)).Normalize();
            }
            if (rightKeyDown)
            {
                cameraForward =
                    cameraForward.Sub(Helper.Cross(down, cameraForward).Normalize().Times(0.1f).Times(scale)).Normalize();
            }
            if (pageDownKeyDown)
            {
                cameraPosition = cameraPosition.Sub(new Float4(0, 1.0f, 0, 0.0f));
                cameraForward = cameraForward.Add(down.Times(-0.2f)).Normalize();
            }
            if (pageUpKeyDown)
            {
                cameraPosition = cameraPosition.Add(new Float4(0, 1.0f, 0, 0.0f));
                cameraForward = cameraForward.Sub(down.Times(-0.2f)).Normalize();    
            }

            Redraw();
        }
    }

    public class DrawPanel : Panel
    {
        public DrawPanel()
            : base()
        {
            //The trick here to avoid "flickering": We enable double buffering
            //but this enabling is only possible in derived panels because the DoubleBuffer
            //is protected
            SetStyle(ControlStyles.OptimizedDoubleBuffer, true);
            SetStyle(ControlStyles.UserPaint, true);
            SetStyle(ControlStyles.AllPaintingInWmPaint, true);            
            UpdateStyles();
        }
    }

    [ComVisible(false), SuppressUnmanagedCodeSecurity()]
    internal class NativeMethods
    {
        [DllImport("user32.dll", CharSet = CharSet.Auto, ExactSpelling = true, CallingConvention = CallingConvention.Winapi)]
        public static extern short GetKeyState(int keyCode);

        public static int HIWORD(int n)
        {
            return ((n >> 16) & 0xffff/*=~0x0000*/);
        }

        public static int LOWORD(int n)
        {
            return (n & 0xffff/*=~0x0000*/);
        }
    } 
}
